package com.example.leetcode.graph.common.unionfind.onedimension.pro.unionbyrank;

/**
 * 按秩合并优化后的并查集
 */
public class UnionFind {
    int[] root;
    // 这里的「秩」指的是每个顶点所处的高度
    int[] rank;

    public UnionFind(int size) {
        root = new int[size];
        rank = new int[size];
        for (int i = 0; i < size; i++) {
            root[i] = i;
            // 初始秩都为1
            rank[i] = 1;
        }
    }

    public int find(int x) {
        while (x != root[x]) {
            x = root[x];
        }
        return x;
    }

    /**
     * 这里的「秩」指的是每个顶点所处的高度。
     * 我们每次 union 两个顶点的时候，选择根节点的时候不是随机的选择某个顶点的根节点，
     * 而是将「秩」大的那个根节点作为两个顶点的根节点，
     * 换句话说，我们将低的树合并到高的树之下，将高的树的根节点作为两个顶点的根节点。
     * 这样，我们就避免了所有的顶点连成一条线，这就是按秩合并优化的「并查集」。
     *
     * @param x
     * @param y
     */
    public void union(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);
        if (rootX != rootY) {
            if (rank[rootX] > rank[rootY]) {
                root[rootY] = rootX;
            } else if (rank[rootX] < rank[rootY]) {
                root[rootX] = rootY;
            } else {
                // 假设原先X和Y的树高都是h，则合并后，Y成了X的子树，X的树高成了h + 1
                // 将Y合并到X下，新的树根是X
                root[rootY] = rootX;
                // X树高因为多了一个子树Y，而子树Y的高度与原先X的高度h一致，所以X新的树高比原先多了1
                rank[rootX] += 1;
            }
        }
    }

    public boolean connected(int x, int y) {
        return find(x) == find(y);
    }
}

